
# -*- coding: utf-8 -*-
# coding=utf-8
# coding: utf-8

import cv2
import easyocr
import numpy as np
from PIL import Image

import matplotlib
matplotlib.use('TkAgg')  # 或尝试 'Qt5Agg'、'WXAgg'
import matplotlib.pyplot as plt

import datetime as dt

import sys,io,os
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')


def getAppendPath(src:str,append:str)->str:
    arr=src.split('.')
    return f'{arr[0]}_{append}.{arr[1]}'

def writeLog(fn:str,memo:str):
    with open(fn,'a',encoding='utf-8') as f:
        f.write(f'{memo}\n')


class PlateRecognizer:
    def __init__(self):
        # 初始化EasyOCR阅读器，支持中文和英文
        self.reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)
        
    def preprocess_image(self, image_path):
        """图像预处理：灰度化、直方图均衡化等操作"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法读取图像文件")
            
        # 转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # 直方图均衡化增强对比度
        equalized = cv2.equalizeHist(gray)
        
        return img, gray, equalized
    
    def locate_plate(self, image):
        """车牌定位：使用多种方法定位车牌区域"""
        # 方法1：颜色定位（蓝色车牌）
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 蓝色车牌范围
        lower_blue = np.array([100, 50, 50])
        upper_blue = np.array([140, 255, 255])
        blue_mask = cv2.inRange(hsv, lower_blue, upper_blue)
        
        # 方法2：边缘检测
        edges = cv2.Canny(image, 80, 200)

        kernel = np.ones((5, 19), np.uint8)
        closing = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
        opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
        
        # 寻找轮廓
        contours, _ = cv2.findContours(opening, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        plate_candidates = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            
            # 筛选可能为车牌的轮廓
            aspect_ratio = w / h
            if 2.0 < aspect_ratio < 5.0 and w > 100 and h > 30:
                plate_candidates.append((x, y, w, h))
                
        return plate_candidates
    
    def recognize_plate(self, image_path):
        """主识别函数：检测并识别车牌"""
        try:
            # 图像预处理
            original, gray, enhanced = self.preprocess_image(image_path)
            
            # 车牌定位
            plate_regions = self.locate_plate(original)

            x,y,w,h=plate_regions[0]
            aoi=getAppendPath(image_path,'aoi')
            # cv2.imwrite(aoi,original[y:y+h,x:x+w])
            
            results = []
            for i, (x, y, w, h) in enumerate(plate_regions):
                # 提取车牌区域
                plate_roi = original[y:y+h, x:x+w]
                
                # 使用EasyOCR识别文字
                ocr_result = self.reader.readtext(plate_roi)
                
                if ocr_result:
                    text = ocr_result[0][1]
                    confidence = ocr_result[0][2]
                    
                    results.append({
                        'plate_text': text,
                        'confidence': confidence,
                        'coordinates': (x, y, w, h)
                    })
            
            return results
            
        except Exception as e:
            print(f"识别过程中出错: {e}")
            return []
        

def main():
    t1=dt.datetime.now()
    # 创建识别器实例
    recognizer = PlateRecognizer()
    pic_ext=['jpg','png','bmp']
    dir=os.path.dirname(sys.argv[0])
    os.chdir('cp_img')
    for f in  os.listdir():
        arrf= str(f).split('.')
        if len(arrf)==2:
            if pic_ext.__contains__(arrf[1]):
                test_image = f'{dir}/cp_img/{f}'

                # print(test_image) 
                
                t1=dt.datetime.now()
                # 执行识别
                results = recognizer.recognize_plate(test_image)
                t2=dt.datetime.now()
    
                
                # 输出结果
                temp=[]
                if results:
                    temp.append("识别结果:")
                    for i, result in enumerate(results):
                        temp.append(f"车牌 {i+1}: {result['plate_text']}")
                        temp.append(f"置信度: {result['confidence']:.2f}")
                        temp.append(f"位置: {result['coordinates']}")
                        temp.append(f'耗时(ms):{(t2-t1).microseconds/1000}')
                        temp.append("-" * 30)
                else:
                    temp.append("未检测到车牌")
                for i in temp:
                    writeLog(f'{dir}/log/{t1.year}-{t1.month}-{t1.day}_cpsb.txt',i)
                    print(i)
                    pass

if __name__ == "__main__":
    
    main()

